import socket
import json
import struct
import time
from recommend_trie_model import TrieNode, load_keywords, model, find_similar_keywords  # 导入所需功能

# 初始化 SentenceModel 已在子模块中初始化
# model_path = "model/shibing624-text2vec-base-chinese"
# model = SentenceModel(model_path)

# 初始化 Trie 和加载词典
# keywords_file_path = "dict/merge_dict.dat"  # 小词典
# keywords_file_path = "dict/merged_dictionary.dat" # 大词典
# keywords_dict = load_keywords(keywords_file_path)  # 加载关键词

# trie = TrieNode()  # 初始化 Trie
# for word in keywords_dict:
    # trie.insert(word)  # 插入关键词到 Trie 中

# 编码词典中的关键词
# keyword_embeddings = model.encode(keywords_dict)

def receive_message(conn):
    data = b''
    try:
        while True:
            part = conn.recv(1)
            if not part:
                break  # 客户端已关闭连接
            data += part
            if data.endswith(b'\n'):
                break
    except ConnectionResetError:
        print("Connection was reset by the client.")
    except socket.error as e:
        print(f"Socket error: {e}")
    return data.strip()

def send_message(conn, message):
    try:
        encoded_message = message.encode()
        message_length = len(encoded_message)
        conn.sendall(struct.pack('Q', message_length))
        conn.sendall(encoded_message)
    except ConnectionResetError:
        print("Connection was reset by the client while sending.")
    except socket.error as e:
        print(f"Error sending data: {e}")

def handle_query(data):
    try:
        request = json.loads(data.decode())
        if request["query_id"] == 1:
            query = request["msg"]
            similar_keywords, elapsed_time = find_similar_keywords(query) # 使用新的查找函数
            print(f"与 '{query}' 语义上接近的关键词和相应的相似度分数为:")
            for keyword, score in similar_keywords:
                print(f"{keyword}: {score:.4f}")
            print(f"查询耗时: {elapsed_time:.4f} 秒\n")

            response = {
                "msgID": 1,
                "msg": [kw for kw, _ in similar_keywords]
            }
        else:
            response = {"msgID": 0, "error": "Invalid query_id"}
        return json.dumps(response)
    except json.JSONDecodeError:
        return json.dumps({"msgID": 0, "error": "Invalid JSON format"})

def start_server(host='127.0.0.1', port=8888):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind((host, port))
        s.listen()
        print(f"Server started on {host}:{port}")
        while True:
            try:
                conn, addr = s.accept()
                with conn:
                    print(f"Connected by {addr}")
                    while True:
                        data = receive_message(conn)
                        if not data:
                            break
                        response = handle_query(data)
                        send_message(conn, response)
            except socket.error as e:
                print(f"Server socket error: {e}")

if __name__ == "__main__":
    start_server()
